from abc import ABC, abstractmethod
import numpy as np
from typing import Union
from light_rl.model.model import Model
from light_rl.utils.buffer import ReplayBuffer
from light_rl.utils.utils import load_config
from light_rl.utils.utils import get_device
from light_rl.utils.dict_wrapper import DictWrapper
import os
import json
from torch.utils.tensorboard import SummaryWriter


class AgentBase(ABC):
    def __init__(self, config_path: str, writer: SummaryWriter = None, model_dir: str = "") -> None:

        self.full_config = load_config(config_path)
        self.config = self.full_config.agent

        self.device = get_device(self.full_config.device)
        self.resume = self.full_config.resume

        self.writer = writer

        self.model_dir = model_dir
        self.step = 0

    @abstractmethod
    def react(self, state: np.ndarray, train: bool = True) -> Union[np.ndarray, int]:
        """
        根据当前状态决定采取什么样的动作
        :param state: np.ndarray(env_num, state_dim) 状态,传入一个
        :param train: bool 是否用于训练，如果用于训练，可能需要探索，不用于训练不用探索
        :return: np.ndarray(env_num, -1) 返回action，用于与环境交互,如果是离散动作，shape[1]=1,如果是连续动作shape[1]=action_dim
        """
        raise NotImplementedError

    @abstractmethod
    def learn(self, state: np.ndarray, action: Union[np.ndarray, int], reward, next_state: np.ndarray,
              done: bool) -> None:
        """
        根据经验进行强化学习算法的学习，更新网络参数
        :param state: np.ndarray(state_dim) 状态
        :param action: np.ndarray(action_dim)|int 动作，如果是离散动作 int,如果是连续动作action_dim
        :param reward: float 奖励
        :param next_state: np.ndarray(state_dim) 下一个状态
        :param done: bool 是否结束
        :return: None
        """
        self.step += 1
        self.auto_save()

    @staticmethod
    def soft_update(target: Model, source: Model, tau: float):
        for target_param, source_param in zip(target.parameters(), source.parameters()):
            target_param.data.copy_(
                tau * source_param.data + (1 - tau) * target_param.data
            )

    def resume_model(self):
        """
        从模型文件夹中恢复模型
        :return: None
        """
        if not self.resume:
            return
        if self.model_dir == "":
            return
        model_dir = os.path.dirname(self.model_dir)
        if not os.path.exists(model_dir):
            os.makedirs(model_dir)
            return
        models = os.listdir(model_dir)
        models.sort()
        if len(models) == 0:
            print("No model to load")
            return

        if "load_dir" not in self.full_config or self.full_config.load_dir == "":
            model_path = os.path.join(model_dir, models[-1])
        else:
            if self.full_config.load_dir not in models:
                model_path = os.path.join(model_dir, models[-1])
            else:
                model_path = os.path.join(model_dir, self.full_config.load_dir)

        print("load model from {}".format(model_path))

        checkpoint_list = [int(i) for i in os.listdir(model_path)]
        checkpoint_list.sort()
        if len(checkpoint_list) == 0:
            print("No checkpoint to load")
            return
        if "checkpoint" not in self.full_config or self.full_config.checkpoint == -1:
            checkpoint_path = os.path.join(model_path, str(checkpoint_list[-1]))
        else:
            if self.full_config.checkpoint not in checkpoint_list:
                checkpoint_path = os.path.join(model_path, str(checkpoint_list[-1]))
            else:
                checkpoint_path = os.path.join(model_path, str(self.full_config.checkpoint))

        print("load checkpoint from {}".format(checkpoint_path))
        self.load(checkpoint_path)

    def auto_save(self):
        """
        自动保存模型
        :return: None
        """
        if self.model_dir == "":
            return
        if self.full_config.auto_save_interval == -1:
            return
        # print("step: {} , checkpoint: {}, mod: {}".format(self.step, self.full_config.auto_save_interval * 1000, self.step % (self.full_config.auto_save_interval * 1000)))
        self.checkpoint = self.step // 1000
        if self.step % (self.full_config.auto_save_interval * 1000) == 0:
            checkpoint_path = os.path.join(self.model_dir, str(self.checkpoint))
            self.save(checkpoint_path)
            print("checkpoint {} save model to {}".format(str(self.checkpoint), checkpoint_path))

    def end_save(self):
        """
        保存最终模型
        :return: None
        """
        if self.model_dir == "":
            return
        self.checkpoint = self.step // 1000
        checkpoint_path = os.path.join(self.model_dir,  str(self.checkpoint))
        self.save(checkpoint_path)
        print("end save model to {}".format(checkpoint_path))

    def save(self, path: str) -> None:
        """
        保存智能体的数据，包括模型和离线的经验，在下一次训练的时候可以加载模型继续训练，或者加载模型用于测试或部署
        :param path: str 保存的路径,一般要求传入一个文件夹路径，然后在这个文件夹下面保存模型和经验，如果文件夹不存在，会自动创建
        :return: None
        """

        os.makedirs(path, exist_ok=True)
        for attr_name, attr_value in self.__dict__.items():
            if isinstance(attr_value, Model):
                attr_value.save(os.path.join(path, attr_name + ".pth"))
            if isinstance(attr_value, ReplayBuffer):
                attr_value.save(os.path.join(path, "buffer.pkl"))

        config_path = os.path.join(path, "config.json")
        json.dump(self.config, open(config_path, "w"))

    def load(self, path: str) -> bool:
        """
        加载训练的模型和经验，如果加载成功，返回True，否则返回False
        :param path: str 保存的路径,一般要求传入一个文件夹路径，然后在这个文件夹下面保存模型和经验
        :return: bool
        """
        if not os.path.exists(path):
            return False
        for attr_name, attr_value in self.__dict__.items():
            if isinstance(attr_value, Model):
                attr_value.load(os.path.join(path, attr_name + ".pth"))
            if isinstance(attr_value, ReplayBuffer):
                attr_value.save(os.path.join(path, "buffer.pkl"))

        config_path = os.path.join(path, "config.json")
        if os.path.exists(config_path):
            self.config = DictWrapper(json.load(open(config_path, "r")))
        return True
